package com.zzh.partnersys.ai.util;

import org.apache.commons.math3.linear.EigenDecomposition;
import org.apache.commons.math3.linear.MatrixUtils;
import org.apache.commons.math3.linear.RealMatrix;
import org.apache.commons.math3.linear.RealVector;
import org.apache.commons.math3.stat.correlation.Covariance;

import java.util.ArrayList;
import java.util.List;

public class DimensionReducerUtil {
    // 目标维度（768）
    private static final int TARGET_DIM = 768;
    // 源维度（1536，和 Spring AI 生成的向量维度一致）
    private static final int SOURCE_DIM = 1536;

    public static float[] reduceDims(float[] original, int targetDims) {
        int originalDims = original.length;
        float[] reduced = new float[targetDims];
        // 按比例合并维度（示例逻辑，需替换为标准 PCA 降维）
        for (int i = 0; i < targetDims; i++) {
            int start = (i * originalDims) / targetDims;
            int end = ((i + 1) * originalDims) / targetDims;
            float sum = 0;
            for (int j = start; j < end; j++) {
                sum += original[j];
            }
            reduced[i] = sum / (end - start);
        }
        return reduced;
    }


    /**
     * 批量将 1536 维向量降维到 768 维（修复类型转换和语法问题）
     * @param highDimVectors 1536 维向量列表（每个向量长度必须是 1536）
     * @return 768 维向量列表
     * @throws IllegalArgumentException 输入向量维度错误或样本量不足
     */
    public static List<double[]> pcaReduce(List<float[]> highDimVectors) {
        // 前置校验：避免无效输入
        if (highDimVectors == null || highDimVectors.isEmpty()) {
            throw new IllegalArgumentException("高维向量列表不能为空");
        }
        // 校验每个向量的维度是否为 1536
        for (int i = 0; i < highDimVectors.size(); i++) {
            float[] vector = highDimVectors.get(i);
            if (vector == null || vector.length != SOURCE_DIM) {
                throw new IllegalArgumentException("第 " + (i + 1) + " 个向量维度错误，必须是 " + SOURCE_DIM + " 维");
            }
        }

        // 1. 转换为矩阵（行：样本数，列：1536 维）→ 核心修复：float[] 转 double[]
        int sampleCount = highDimVectors.size();
        double[][] data = new double[sampleCount][SOURCE_DIM];
        for (int i = 0; i < sampleCount; i++) {
            float[] floatVector = highDimVectors.get(i);
            // 手动将 float[] 转为 double[]（解决 createRealVector 不支持 float[] 的问题）
            double[] doubleVector = new double[SOURCE_DIM];
            for (int j = 0; j < SOURCE_DIM; j++) {
                doubleVector[j] = floatVector[j]; // 隐式转换（float → double 无精度损失）
            }
            data[i] = doubleVector;
        }
        RealMatrix matrix = MatrixUtils.createRealMatrix(data);

        // 2. 计算协方差矩阵（PCA 核心步骤）
        Covariance covariance;
        try {
            covariance = new Covariance(matrix);
        } catch (IllegalArgumentException e) {
            throw new IllegalArgumentException("PCA 降维需要至少 2 个样本（当前样本数：" + sampleCount + "）", e);
        }
        RealMatrix covMatrix = covariance.getCovarianceMatrix();

        // 3. 特征值分解，获取前 768 个主成分（保留最多信息）
        EigenDecomposition eigenDecomp = new EigenDecomposition(covMatrix);
        RealMatrix eigenVectors = eigenDecomp.getV(); // 特征向量矩阵（1536 x 1536）
        // 截取前 768 个特征向量（列维度：0 到 767）→ 核心修复：索引范围正确
        RealMatrix topEigenVectors = eigenVectors.getSubMatrix(
                0, SOURCE_DIM - 1, // 行范围：0-1535（1536 维）
                0, TARGET_DIM - 1  // 列范围：0-767（取前 768 个主成分）
        );

        // 4. 矩阵乘法：1536 维 → 768 维（样本数 x 1536）×（1536 x 768）=（样本数 x 768）
        RealMatrix lowDimMatrix = matrix.multiply(topEigenVectors);

        // 5. 转换为 float[] 列表返回 → 核心修复：用普通循环替代 Lambda，避免 stream() 语法问题
        List<double[]> result = new ArrayList<>(sampleCount);
        double[][] lowDimData = lowDimMatrix.getData(); // 行：样本数，列：768
        for (double[] row : lowDimData) {
            double[] floatRow = new double[TARGET_DIM];
            // double → float 无精度损失（工程上可接受）
            System.arraycopy(row, 0, floatRow, 0, TARGET_DIM);
            result.add(floatRow);
        }

        return result;
    }



    /**
     * 平均池化降维：1536维 → 768维（每组2个元素取平均，无收敛问题）
     * @param highDimVectors 1536维向量列表
     * @return 768维向量列表
     */
    public static List<double[]> averagePoolReduce(List<float[]> highDimVectors) {
        // 前置校验
        if (highDimVectors == null || highDimVectors.isEmpty()) {
            throw new IllegalArgumentException("高维向量列表不能为空");
        }
        for (int i = 0; i < highDimVectors.size(); i++) {
            float[] vector = highDimVectors.get(i);
            if (vector == null || vector.length != SOURCE_DIM) {
                throw new IllegalArgumentException("第 " + (i + 1) + " 个向量维度错误，必须是 " + SOURCE_DIM + " 维");
            }
        }

        List<double[]> result = new ArrayList<>(highDimVectors.size());
        for (float[] highDimVec : highDimVectors) {
            double[] lowDimVec = new double[TARGET_DIM];
            // 每2个元素取平均，压缩到768维
            for (int i = 0; i < TARGET_DIM; i++) {
                // 第i组 = 第2i个元素 + 第2i+1个元素 的平均值
                lowDimVec[i] = (highDimVec[2 * i] + highDimVec[2 * i + 1]) / 2.0d;
            }
            result.add(lowDimVec);
        }
        return result;
    }

}